import pandas as pd
import numpy as np
from sklearn.preprocessing import MinMaxScaler
import os

class StockDividendPreprocessor:
    """
    股票分红数据预处理类
    用于对爬取的股票分红数据进行清洗、分析和预处理
    """
    def __init__(self):
        self.data = None
        
    def load_data(self, file_path):
        """
        从CSV文件加载数据
        """
        try:
            # 检查文件是否存在
            if not os.path.exists(file_path):
                print(f"错误：文件 {file_path} 不存在！")
                return False
            
            # 从CSV文件加载数据
            self.data = pd.read_csv(file_path, encoding='utf-8-sig')
            print(f"成功加载数据，文件路径：{file_path}")
            return True
        except Exception as e:
            print(f"加载数据时发生错误：{str(e)}")
            return False
    
    def check_data_shape(self):
        """
        检查并输出数据形状
        """
        if self.data is not None:
            rows, cols = self.data.shape
            print(f"数据形状：{rows}行 × {cols}列")
            print(f"列名：{list(self.data.columns)}")
            return rows, cols
        else:
            print("数据尚未加载")
            return None, None
    
    def check_duplicates(self, columns=['代码', '名称', '详细']):
        """
        检查指定列的重复值
        """
        if self.data is not None:
            print("\n重复值检查结果：")
            for col in columns:
                if col in self.data.columns:
                    duplicate_count = self.data[col].duplicated().sum()
                    print(f"列 '{col}' 中的重复值数量：{duplicate_count}")
                    # 如果有重复值，显示重复的具体值
                    if duplicate_count > 0:
                        duplicates = self.data[self.data[col].duplicated()][col].unique()
                        print(f"列 '{col}' 中重复的值示例：{list(duplicates)[:5]}...")  # 只显示前5个重复值
                else:
                    print(f"列 '{col}' 不存在于数据中")
        else:
            print("数据尚未加载")
    
    def show_first_rows(self, n=5):
        """
        展示数据的前n行
        """
        if self.data is not None:
            print(f"\n数据的前{n}行：")
            print(self.data.head(n))
            return self.data.head(n)
        else:
            print("数据尚未加载")
            return None
            
    def adjust_data_issues(self):
        """
        观察数据存在的问题并进行调整
        """
        if self.data is None:
            print("错误：请先加载数据！")
            return False
        
        try:
            print("\n开始调整数据...")
            
            # 1. 处理缺失值
            original_missing = self.data.isnull().sum().sum()
            # 对不同类型的列采取不同的缺失值处理策略
            for col in self.data.columns:
                if col in ['代码', '名称']:
                    # 对于关键标识列，填充默认值
                    self.data[col] = self.data[col].fillna('未知')
                elif col in ['累计股息(%)', '年均股息(%)', '分红次数', '融资总额(亿元)']:
                    # 对于数值型列，填充0或均值
                    if self.data[col].dtype in ['float64', 'int64']:
                        self.data[col] = self.data[col].fillna(self.data[col].mean())
                    else:
                        # 尝试转换为数值类型
                        try:
                            self.data[col] = pd.to_numeric(self.data[col], errors='coerce')
                            self.data[col] = self.data[col].fillna(self.data[col].mean())
                        except:
                            self.data[col] = self.data[col].fillna(0)
                else:
                    # 其他列填充空字符串
                    self.data[col] = self.data[col].fillna('')
            
            adjusted_missing = self.data.isnull().sum().sum()
            print(f"缺失值处理：处理前 {original_missing} 个，处理后 {adjusted_missing} 个")
            
            # 2. 处理重复行
            original_rows = len(self.data)
            self.data = self.data.drop_duplicates()
            adjusted_rows = len(self.data)
            print(f"重复行处理：处理前 {original_rows} 行，处理后 {adjusted_rows} 行")
            
            # 3. 确保代码列格式正确
            if '代码' in self.data.columns:
                # 尝试将代码转换为字符串格式，去除可能的空格
                self.data['代码'] = self.data['代码'].astype(str).str.strip()
                print("代码列格式已调整为字符串类型")
            
            print("数据调整完成")
            return True
        except Exception as e:
            print(f"数据调整时发生错误：{str(e)}")
            return False
    
    def process_listing_date(self):
        """
        处理上市日期，提取年、月、日变量
        """
        if self.data is not None:
            try:
                if '上市日期' not in self.data.columns:
                    print("警告：'上市日期'列不存在于数据中！")
                    # 如果没有上市日期列，尝试从模拟数据的'分红日期'列生成上市日期
                    if '分红日期' in self.data.columns:
                        print("尝试使用'分红日期'列生成上市日期...")
                        # 假设上市日期早于分红日期1-5年
                        np.random.seed(42)  # 设置随机种子，保证结果可复现
                        years_ago = np.random.randint(1, 6, size=len(self.data))
                        # 转换分红日期为日期类型
                        self.data['分红日期'] = pd.to_datetime(self.data['分红日期'], errors='coerce')
                        # 生成上市日期
                        self.data['上市日期'] = self.data['分红日期'].apply(lambda x: x - pd.DateOffset(years=np.random.choice(years_ago)) if pd.notnull(x) else pd.NaT)
                    else:
                        print("无法生成上市日期，跳过此步骤")
                        return False
                
                # 转换为日期类型
                self.data['上市日期'] = pd.to_datetime(self.data['上市日期'], errors='coerce')
                
                # 提取年、月、日
                self.data['上市年份'] = self.data['上市日期'].dt.year
                self.data['上市月份'] = self.data['上市日期'].dt.month
                self.data['上市日'] = self.data['上市日期'].dt.day
                
                print("上市日期处理完成，已提取年、月、日变量")
                return True
            except Exception as e:
                print(f"处理上市日期时发生错误：{str(e)}")
                return False
        else:
            print("数据尚未加载")
            return False
    
    def normalize_numerical_variables(self, columns=['累计股息(%)', '年均股息(%)', '分红次数', '融资总额(亿元)']):
        """
        对数值型变量进行归一化处理
        """
        if self.data is not None:
            try:
                # 创建归一化器
                scaler = MinMaxScaler()
                
                # 准备要归一化的数据
                normalized_data = self.data.copy()
                
                # 存储需要归一化的列
                valid_columns = []
                for col in columns:
                    if col in normalized_data.columns:
                        # 尝试将列转换为数值类型
                        normalized_data[col] = pd.to_numeric(normalized_data[col], errors='coerce')
                        valid_columns.append(col)
                    else:
                        print(f"列 '{col}' 不存在于数据中")
                
                if valid_columns:
                    # 进行归一化
                    # 先填充NaN值为该列的均值
                    for col in valid_columns:
                        mean_value = normalized_data[col].mean()
                        normalized_data[col] = normalized_data[col].fillna(mean_value)
                        
                    # 执行归一化
                    normalized_values = scaler.fit_transform(normalized_data[valid_columns])
                    
                    # 创建新的列名
                    normalized_columns = [f'{col}_归一化' for col in valid_columns]
                    
                    # 将归一化后的数据添加到原数据中
                    for i, col in enumerate(normalized_columns):
                        self.data[col] = normalized_values[:, i]
                    
                    print(f"归一化处理完成，已添加以下归一化列：{normalized_columns}")
                else:
                    print("没有找到有效的数值型列进行归一化")
                
                return True
            except Exception as e:
                print(f"归一化处理时发生错误：{str(e)}")
                return False
        else:
            print("数据尚未加载")
            return False
    
    def save_processed_data(self, output_file='新浪财经历史分红数据_预处理.csv'):
        """
        保存预处理后的数据到CSV文件
        """
        if self.data is not None:
            try:
                # 确保输出目录存在
                output_dir = os.path.dirname(output_file)
                if output_dir and not os.path.exists(output_dir):
                    os.makedirs(output_dir)
                    
                # 保存到CSV文件
                self.data.to_csv(output_file, index=False, encoding='utf-8-sig')
                print(f"预处理后的数据已成功保存到：{output_file}")
                return True
            except Exception as e:
                print(f"保存预处理后的数据时发生错误：{str(e)}")
                return False
        else:
            print("没有可保存的预处理数据")
            return False
    
    def run_complete_preprocessing(self, input_file, output_file):
        """
        运行完整的预处理流程
        """
        print("开始数据预处理流程...")
        
        # 1. 加载数据
        if not self.load_data(input_file):
            return False
        
        # 2. 检查数据形状
        self.check_data_shape()
        
        # 3. 检查重复值
        self.check_duplicates()
        
        # 4. 展示前5行数据
        self.show_first_rows()
        
        # 5. 调整数据存在的问题
        if not self.adjust_data_issues():
            return False
        
        # 6. 处理上市日期
        self.process_listing_date()
        
        # 7. 归一化数值型变量
        self.normalize_numerical_variables()
        
        # 8. 保存预处理后的数据
        result = self.save_processed_data(output_file)
        
        print("\n数据预处理流程结束")
        return result

if __name__ == '__main__':
    # 创建预处理实例
    preprocessor = StockDividendPreprocessor()
    
    # 运行完整的预处理流程
    # 注意：在运行预处理脚本之前，请确保已经成功运行了爬虫脚本获取数据
    preprocessor.run_complete_preprocessing(
        input_file='e:/zhang-zhans-homework/chapter3/新浪财经历史分红数据.csv',
        output_file='e:/zhang-zhans-homework/chapter3/新浪财经历史分红数据_预处理.csv'
    )